from abc import ABCMeta, abstractmethod
from typing import List, Dict
from numpy.typing import NDArray

from hj_dataset_devkit.common_util import check_is_3d_transform_matrix
from hj_dataset_devkit import ObstacleCategory, BoundingBox

from model_frame import DetectObject

__all__ = ['FrameInfo', 'TrackObject', 'ObjTracker']

class FrameInfo():
    def __init__(self, timestamp: int, pose: NDArray, abs_index: int,
                 detect: List[DetectObject]) -> None:
        self.__timestamp = timestamp
        assert check_is_3d_transform_matrix(pose)
        self.__pose = pose
        self.__abs_index = abs_index
        self.__detect = detect

    @property
    def timestamp(self) -> int:
        return self.__timestamp

    @property
    def pose(self) -> NDArray:
        return self.__pose

    @property
    def abs_index(self) -> int:
        return self.__abs_index
    
    @property
    def detect(self) -> List[DetectObject]:
        return self.__detect

class TrackObject:
    def __init__(self, id_: int, category: ObstacleCategory, bbox: BoundingBox,
                 score: float) -> None:
        self.id = id_
        self.category = category
        self.bbox = bbox
        self.score = score

    def to_dict(self) -> dict:
        res = {
            'id': self.id,
            'category': self.category.value,
            'score': self.score
        }
        res.update(self.bbox.to_dict())
        return res

class ObjTracker(metaclass=ABCMeta):
    @abstractmethod
    def __init__(self) -> None:
        pass

    @abstractmethod
    def process(self, frames_info: List[FrameInfo]) -> Dict[int, List[TrackObject]]:
        pass
